#!/usr/bin/env python
#
# Author: Qiming Sun <osirpt.sun@gmail.com>
#

import numpy, time, h5py
from pyscf import lib, symm, df
from pyscf_ccso import dfoutcore, cholesky_eri

dirprd = numpy.array(((0, 1, 2, 3, 4, 5, 6, 7),
          (1, 0, 3, 2, 5, 4, 7, 6),
          (2, 3, 0, 1, 6, 7, 4, 5),
          (3, 2, 1, 0, 7, 6, 5, 4),
          (4, 5, 6, 7, 0, 1, 2, 3),
          (5, 4, 7, 6, 1, 0, 3, 2),
          (6, 7, 4, 5, 2, 3, 0, 1),
          (7, 6, 5, 4, 3, 2, 1, 0)))
   
def ri_blocks(myhf, frozen, filenm='cd_b', auxbas=None):
    mol = myhf.mol
    orb = myhf.mo_coeff[:, frozen:]
    occidx = myhf.mo_occ[frozen:] == 2
    vrtidx = myhf.mo_occ[frozen:] == 0
    norb = len(occidx)
    
    if not mol.symmetry:
        symid2 = numpy.zeros(norb, dtype=int)
    else:
        symid2 = symm.label_orb_symm(mol, mol.irrep_id, mol.symm_orb, orb)
        symid2 = symid2 % 10

    if auxbas == None and myhf.with_df.auxbasis == None:
         print('auxiliary basis set not provided')
         print('It can be provided either with myhf or with ccso_scu!')
    auxbas1 = auxbas
    if auxbas1 == None:
       auxbas1 = myhf.with_df.auxbasis

#   df.outcore.cholesky_eri_b(mol, 'aocderi.dat', auxbasis=auxbas1, dataname='eri_mo')
    aux_nsym, erfile = cholesky_eri.cholesky_eri_b(mol, 'aocderi.dat', auxbasis=auxbas1, dataname='eri_mo')
    
    nirrep = numpy.amax(symid2) + 1
#   print('nirrep:', nirrep)
#   print('symid2:', symid2)
    pop = [0]*nirrep
    vrt = [0]*nirrep
    idxocc = []
    idxvrt = []

    sum_aux = numpy.zeros(nirrep+1, dtype=int)
    for irrep in range(1,nirrep+1):
        sum_aux[irrep] = sum_aux[irrep-1] + aux_nsym[irrep-1]

    naux0 = sum_aux[nirrep]
#   aux_nsym[:] = naux0
    print('aux_nsym', aux_nsym)
    print('sum_aux ', sum_aux)
    for irrep in range(nirrep):
        occir = occidx & (symid2 == irrep)
        vrtir = vrtidx & (symid2 == irrep)
        idxocc.append(occir)
        idxvrt.append(vrtir)
        pop[irrep] = numpy.count_nonzero(occir)
        vrt[irrep] = numpy.count_nonzero(vrtir)
        print('irrep, occ, vrt', irrep, pop[irrep], vrt[irrep])
        for jrrep in range(nirrep):
            occjr = occidx & (symid2 == jrrep)
            vrtjr = vrtidx & (symid2 == jrrep)
            krrep = dirprd[irrep, jrrep]
            idxrange = (sum_aux[krrep], sum_aux[krrep+1])
#           idxrange = (0, naux0)

            strnd = str(irrep) + str(jrrep)
            if numpy.count_nonzero(occir) !=0 and numpy.count_nonzero(occjr) !=0:
               datanm = 'eroo' + strnd
               naux = dfoutcore.general(mol, (orb[:,occjr], orb[:,occir]), idxrange, filenm, datanm,
                                        tmperi=erfile,  compact=False)
#              print('datanm:', datanm)
#              print('sizej, sizei', orb[:,occjr].shape[1], orb[:,occir].shape[1])
#              feri = h5py.File('cd_b', 'r')
#              print('shape of feri', feri[datanm].shape)
#              feri.close()
            if numpy.count_nonzero(vrtir) !=0 and numpy.count_nonzero(vrtjr) !=0:
               datanm = 'ervv' + strnd
               naux = dfoutcore.general(mol, (orb[:,vrtjr], orb[:,vrtir]), idxrange, filenm, datanm,
                                        tmperi=erfile,  compact=False)
#              print('datanm:', datanm)
#              print('sizej, sizei', orb[:,vrtjr].shape[1], orb[:,vrtir].shape[1])
#              feri = h5py.File('cd_b', 'r')
#              print('shape of feri', feri[datanm].shape)
#              feri.close()
            if numpy.count_nonzero(vrtir) !=0 and numpy.count_nonzero(occjr) !=0:
               datanm = 'ervo' + strnd
               naux = dfoutcore.general(mol, (orb[:,occjr], orb[:,vrtir]), idxrange, filenm, datanm,
                                        tmperi=erfile,  compact=False)
#              print('datanm:', datanm)
#              print('sizej, sizei', orb[:,occjr].shape[1], orb[:,vrtir].shape[1])
#              feri = h5py.File('cd_b', 'r')
#              print('shape of feri', feri[datanm].shape)
#              feri.close()

    return aux_nsym

def uri_blocks(myhf, frozen, filenm_a='cd_a',filenm_b='cd_b' ,auxbas=None):
    mol = myhf.mol
    orb_a = myhf.mo_coeff[0][:, frozen:]
    occidx_a = myhf.mo_occ[0][frozen:] == 1
    vrtidx_a = myhf.mo_occ[0][frozen:] == 0
    norb_a = len(occidx_a)
    print('test uri_blocks')
    if not mol.symmetry:
        symid2_a = numpy.zeros(norb, dtype=int)
    else:
        symid2_a = symm.label_orb_symm(mol, mol.irrep_id, mol.symm_orb, orb_a)
        symid2_a = symid2_a % 10
#    print(mol.symm_orb)

    if auxbas == None and myhf.with_df.auxbasis == None:
         print('auxiliary basis set not provided')
         print('It can be provided either with myhf or with ccso_scu!')
    auxbas1 = auxbas
    if auxbas1 == None:
       auxbas1 = myhf.with_df.auxbasis

#   df.outcore.cholesky_eri_b(mol, 'aocderi.dat', auxbasis=auxbas1, dataname='eri_mo')
    aux_nsym, erfile = cholesky_eri.cholesky_eri_b(mol, 'aocderi.dat', auxbasis=auxbas1, dataname='eri_mo')
    print('aux_nsym', aux_nsym, erfile)
    nirrep = numpy.amax(symid2_a) + 1
#   print('nirrep:', nirrep)
#   print('symid2:', symid2)
    pop_a = [0]*nirrep
    vrt_a = [0]*nirrep
    idxocc_a = []
    idxvrt_a = []

    sum_aux = numpy.zeros(nirrep+1, dtype=int)
    for irrep in range(1,nirrep+1):
        sum_aux[irrep] = sum_aux[irrep-1] + aux_nsym[irrep-1]

    naux0 = sum_aux[nirrep]
#   aux_nsym[:] = naux0
    print('aux_nsym', aux_nsym)
    print('sum_aux ', sum_aux)
    for irrep in range(nirrep):
        occir_a = occidx_a & (symid2_a == irrep)
        vrtir_a = vrtidx_a & (symid2_a == irrep)
        idxocc_a.append(occir_a)
        idxvrt_a.append(vrtir_a)
        pop_a[irrep] = numpy.count_nonzero(occir_a)
        vrt_a[irrep] = numpy.count_nonzero(vrtir_a)
        print('irrep, occ_a, vrt_a', irrep, pop_a[irrep], vrt_a[irrep])
        for jrrep in range(nirrep):
            occjr_a = occidx_a & (symid2_a == jrrep)
            vrtjr_a = vrtidx_a & (symid2_a == jrrep)
            krrep = dirprd[irrep, jrrep]
            idxrange = (sum_aux[krrep], sum_aux[krrep+1])
#           idxrange = (0, naux0)

            strnd = str(irrep) + str(jrrep)
            if numpy.count_nonzero(occir_a) !=0 and numpy.count_nonzero(occjr_a) !=0:
               datanm = 'eooa' + strnd
               naux = dfoutcore.general(mol, (orb_a[:,occjr_a], orb_a[:,occir_a]), idxrange, filenm_a, datanm,
                                        tmperi=erfile,  compact=False)
            if numpy.count_nonzero(vrtir_a) !=0 and numpy.count_nonzero(vrtjr_a) !=0:
               datanm = 'evva' + strnd
               naux = dfoutcore.general(mol, (orb_a[:,vrtjr_a], orb_a[:,vrtir_a]), idxrange, filenm_a, datanm,
                                        tmperi=erfile,  compact=False)
            if numpy.count_nonzero(vrtir_a) !=0 and numpy.count_nonzero(occjr_a) !=0:
               datanm = 'evoa' + strnd
               naux = dfoutcore.general(mol, (orb_a[:,occjr_a], orb_a[:,vrtir_a]), idxrange, filenm_a, datanm,
                                        tmperi=erfile,  compact=False)
               
    orb_b = myhf.mo_coeff[1][:, frozen:]
    occidx_b = myhf.mo_occ[1][frozen:] == 1
    vrtidx_b = myhf.mo_occ[1][frozen:] == 0
    norb_b = len(occidx_b)
    print('test uri_blocks')
    if not mol.symmetry:
        symid2_b = numpy.zeros(norb, dtype=int)
    else:
        symid2_b = symm.label_orb_symm(mol, mol.irrep_id, mol.symm_orb, orb_b)
        symid2_b = symid2_b % 10
#    print('irrepid',mol.irrep_id)

    if auxbas == None and myhf.with_df.auxbasis == None:
         print('auxiliary basis set not provided')
         print('It can be provided either with myhf or with ccso_scu!')
    auxbas1 = auxbas
    if auxbas1 == None:
       auxbas1 = myhf.with_df.auxbasis

#   df.outcore.cholesky_eri_b(mol, 'aocderi.dat', auxbasis=auxbas1, dataname='eri_mo')
    aux_nsym, erfile = cholesky_eri.cholesky_eri_b(mol, 'aocderi.dat', auxbasis=auxbas1, dataname='eri_mo')
    print('aux_nsym', aux_nsym, erfile)
    nirrep = numpy.amax(symid2_b) + 1
#   print('nirrep:', nirrep)
#   print('symid2:', symid2)
    pop_b = [0]*nirrep
    vrt_b = [0]*nirrep
    idxocc_b = []
    idxvrt_b = []

    sum_aux = numpy.zeros(nirrep+1, dtype=int)
    for irrep in range(1,nirrep+1):
        sum_aux[irrep] = sum_aux[irrep-1] + aux_nsym[irrep-1]

    naux0 = sum_aux[nirrep]
#   aux_nsym[:] = naux0
    print('aux_nsym', aux_nsym)
    print('sum_aux ', sum_aux)
    for irrep in range(nirrep):
        occir_b = occidx_b & (symid2_b == irrep)
        vrtir_b = vrtidx_b & (symid2_b == irrep)
        idxocc_b.append(occir_b)
        idxvrt_b.append(vrtir_b)
        pop_b[irrep] = numpy.count_nonzero(occir_b)
        vrt_b[irrep] = numpy.count_nonzero(vrtir_b)
        print('irrep, occ_b, vrt_b', irrep, pop_b[irrep], vrt_b[irrep])
        for jrrep in range(nirrep):
            occjr_b = occidx_b & (symid2_b == jrrep)
            vrtjr_b = vrtidx_b & (symid2_b == jrrep)
            krrep = dirprd[irrep, jrrep]
            idxrange = (sum_aux[krrep], sum_aux[krrep+1])
#           idxrange = (0, naux0)

            strnd = str(irrep) + str(jrrep)
            if numpy.count_nonzero(occir_b) !=0 and numpy.count_nonzero(occjr_b) !=0:
               datanm = 'eoob' + strnd
               naux = dfoutcore.general(mol, (orb_b[:,occjr_b], orb_b[:,occir_b]), idxrange, filenm_b, datanm,
                                        tmperi=erfile,  compact=False)
            if numpy.count_nonzero(vrtir_b) !=0 and numpy.count_nonzero(vrtjr_b) !=0:
               datanm = 'evvb' + strnd
               naux = dfoutcore.general(mol, (orb_b[:,vrtjr_b], orb_b[:,vrtir_b]), idxrange, filenm_b, datanm,
                                        tmperi=erfile,  compact=False)
            if numpy.count_nonzero(vrtir_b) !=0 and numpy.count_nonzero(occjr_b) !=0:
               datanm = 'evob' + strnd
               naux = dfoutcore.general(mol, (orb_b[:,occjr_b], orb_b[:,vrtir_b]), idxrange, filenm_b, datanm,
                                        tmperi=erfile,  compact=False)
    return aux_nsym


def rori_blocks(myhf, frozen, filenm_a='cd_a',filenm_b='cd_b' ,auxbas=None):
    mol = myhf.mol
    orb_a = myhf.mo_coeff.ro[0, :, frozen:]
    occidx_a = myhf.mo_occ[frozen:] == 2
    occidx_a += myhf.mo_occ[frozen:] == 1
    vrtidx_a = myhf.mo_occ[frozen:] == 0
    norb_a = len(occidx_a)
    print('test uri_blocks')
    if not mol.symmetry:
        symid2_a = numpy.zeros(norb, dtype=int)
    else:
        symid2_a = symm.label_orb_symm(mol, mol.irrep_id, mol.symm_orb, orb_a)
        symid2_a = symid2_a % 10
#    print(mol.symm_orb)

    if auxbas == None and myhf.with_df.auxbasis == None:
         print('auxiliary basis set not provided')
         print('It can be provided either with myhf or with ccso_scu!')
    auxbas1 = auxbas
    if auxbas1 == None:
       auxbas1 = myhf.with_df.auxbasis

#   df.outcore.cholesky_eri_b(mol, 'aocderi.dat', auxbasis=auxbas1, dataname='eri_mo')
    aux_nsym, erfile = cholesky_eri.cholesky_eri_b(mol, 'aocderi.dat', auxbasis=auxbas1, dataname='eri_mo')
    print('aux_nsym', aux_nsym, erfile)
    nirrep = numpy.amax(symid2_a) + 1
#   print('nirrep:', nirrep)
#   print('symid2:', symid2)
    pop_a = [0]*nirrep
    vrt_a = [0]*nirrep
    idxocc_a = []
    idxvrt_a = []

    sum_aux = numpy.zeros(nirrep+1, dtype=int)
    for irrep in range(1,nirrep+1):
        sum_aux[irrep] = sum_aux[irrep-1] + aux_nsym[irrep-1]

    naux0 = sum_aux[nirrep]
#   aux_nsym[:] = naux0
    print('aux_nsym', aux_nsym)
    print('sum_aux ', sum_aux)
    for irrep in range(nirrep):
        occir_a = occidx_a & (symid2_a == irrep)
        vrtir_a = vrtidx_a & (symid2_a == irrep)
        idxocc_a.append(occir_a)
        idxvrt_a.append(vrtir_a)
        pop_a[irrep] = numpy.count_nonzero(occir_a)
        vrt_a[irrep] = numpy.count_nonzero(vrtir_a)
        print('irrep, occ_a, vrt_a', irrep, pop_a[irrep], vrt_a[irrep])
        for jrrep in range(nirrep):
            occjr_a = occidx_a & (symid2_a == jrrep)
            vrtjr_a = vrtidx_a & (symid2_a == jrrep)
            krrep = dirprd[irrep, jrrep]
            idxrange = (sum_aux[krrep], sum_aux[krrep+1])
#           idxrange = (0, naux0)

            strnd = str(irrep) + str(jrrep)
            if numpy.count_nonzero(occir_a) !=0 and numpy.count_nonzero(occjr_a) !=0:
               datanm = 'eooa' + strnd
               naux = dfoutcore.general(mol, (orb_a[:,occjr_a], orb_a[:,occir_a]), idxrange, filenm_a, datanm,
                                        tmperi=erfile,  compact=False)
            if numpy.count_nonzero(vrtir_a) !=0 and numpy.count_nonzero(vrtjr_a) !=0:
               datanm = 'evva' + strnd
               naux = dfoutcore.general(mol, (orb_a[:,vrtjr_a], orb_a[:,vrtir_a]), idxrange, filenm_a, datanm,
                                        tmperi=erfile,  compact=False)
            if numpy.count_nonzero(vrtir_a) !=0 and numpy.count_nonzero(occjr_a) !=0:
               datanm = 'evoa' + strnd
               naux = dfoutcore.general(mol, (orb_a[:,occjr_a], orb_a[:,vrtir_a]), idxrange, filenm_a, datanm,
                                        tmperi=erfile,  compact=False)
               
    orb_b = myhf.mo_coeff.ro[1,:, frozen:]
    occidx_b = myhf.mo_occ[frozen:] == 2
    vrtidx_b = myhf.mo_occ[frozen:] == 1
    vrtidx_b += myhf.mo_occ[frozen:] == 0
    norb_b = len(occidx_b)
    print('test uri_blocks')
    if not mol.symmetry:
        symid2_b = numpy.zeros(norb, dtype=int)
    else:
        symid2_b = symm.label_orb_symm(mol, mol.irrep_id, mol.symm_orb, orb_b)
        symid2_b = symid2_b % 10
#    print('irrepid',mol.irrep_id)

    if auxbas == None and myhf.with_df.auxbasis == None:
         print('auxiliary basis set not provided')
         print('It can be provided either with myhf or with ccso_scu!')
    auxbas1 = auxbas
    if auxbas1 == None:
       auxbas1 = myhf.with_df.auxbasis

#   df.outcore.cholesky_eri_b(mol, 'aocderi.dat', auxbasis=auxbas1, dataname='eri_mo')
    aux_nsym, erfile = cholesky_eri.cholesky_eri_b(mol, 'aocderi.dat', auxbasis=auxbas1, dataname='eri_mo')
    print('aux_nsym', aux_nsym, erfile)
    nirrep = numpy.amax(symid2_b) + 1
#   print('nirrep:', nirrep)
#   print('symid2:', symid2)
    pop_b = [0]*nirrep
    vrt_b = [0]*nirrep
    idxocc_b = []
    idxvrt_b = []

    sum_aux = numpy.zeros(nirrep+1, dtype=int)
    for irrep in range(1,nirrep+1):
        sum_aux[irrep] = sum_aux[irrep-1] + aux_nsym[irrep-1]

    naux0 = sum_aux[nirrep]
#   aux_nsym[:] = naux0
    print('aux_nsym', aux_nsym)
    print('sum_aux ', sum_aux)
    for irrep in range(nirrep):
        occir_b = occidx_b & (symid2_b == irrep)
        vrtir_b = vrtidx_b & (symid2_b == irrep)
        idxocc_b.append(occir_b)
        idxvrt_b.append(vrtir_b)
        pop_b[irrep] = numpy.count_nonzero(occir_b)
        vrt_b[irrep] = numpy.count_nonzero(vrtir_b)
        print('irrep, occ_b, vrt_b', irrep, pop_b[irrep], vrt_b[irrep])
        for jrrep in range(nirrep):
            occjr_b = occidx_b & (symid2_b == jrrep)
            vrtjr_b = vrtidx_b & (symid2_b == jrrep)
            krrep = dirprd[irrep, jrrep]
            idxrange = (sum_aux[krrep], sum_aux[krrep+1])
#           idxrange = (0, naux0)

            strnd = str(irrep) + str(jrrep)
            if numpy.count_nonzero(occir_b) !=0 and numpy.count_nonzero(occjr_b) !=0:
               datanm = 'eoob' + strnd
               naux = dfoutcore.general(mol, (orb_b[:,occjr_b], orb_b[:,occir_b]), idxrange, filenm_b, datanm,
                                        tmperi=erfile,  compact=False)
            if numpy.count_nonzero(vrtir_b) !=0 and numpy.count_nonzero(vrtjr_b) !=0:
               datanm = 'evvb' + strnd
               naux = dfoutcore.general(mol, (orb_b[:,vrtjr_b], orb_b[:,vrtir_b]), idxrange, filenm_b, datanm,
                                        tmperi=erfile,  compact=False)
            if numpy.count_nonzero(vrtir_b) !=0 and numpy.count_nonzero(occjr_b) !=0:
               datanm = 'evob' + strnd
               naux = dfoutcore.general(mol, (orb_b[:,occjr_b], orb_b[:,vrtir_b]), idxrange, filenm_b, datanm,
                                        tmperi=erfile,  compact=False)
    return aux_nsym




